import numpy as np
import matplotlib.pyplot as plt
import os
from matplotlib.lines import lineStyles
from matplotlib.animation import FuncAnimation

root_path = './output'
def load(res_index, res_tag, res_type):
    file = open(root_path+res_index+'/'+res_tag+'_'+res_type+'.res', 'r')
    lines = file.readlines()

    print(root_path+res_index+'/'+res_tag+'_'+res_type+'.res')
    data = []
    for line in lines:
        row = line.split()
        row_data = [float(x) for x in row]
        data.append(row_data)
    t = np.array([row[0] for row in data])
    x = np.array([row[1] for row in data])
    return t,x

def generate(res_index, res_tag, style, color):
    plt.subplot(4, 1, 1)
    plt.title('value')
    t,x = load(res_index, res_tag,'val')
    plt.plot(t,x,linestyle = style, color =color, label = res_tag + '_' + res_index)
    plt.subplot(4, 1, 2)
    plt.title('velocity')
    t, x = load(res_index, res_tag, 'vel')
    plt.plot(t, x, linestyle=style, color=color, label=res_tag + '_' + res_index)
    plt.subplot(4, 1, 3)
    plt.title('frequency')
    t,x = load(res_index, res_tag, 'fre')
    plt.plot(t,x, linestyle=style, color=color, label= res_tag + '_' + res_index)
    plt.subplot(4, 1, 4)
    plt.title('phase')
    t,x = load(res_index, res_tag, 'phase')
    plt.plot(t,x, linestyle=style, color=color, label= res_tag + '_' + res_index)

def generate_DPP(res_index, res_tag, style, color):
    l1 = 1.0
    l2 = 1.0
    m1 = 1.0
    m2 = 1.0
    h = 0.2

    t, t1 = load(res_index, res_tag, 't1')
    t, t2 = load(res_index, res_tag, 't2')
    t, w1 = load(res_index, res_tag, 'w1')
    t, w2 = load(res_index, res_tag, 'w2')

    x1 = l1 * np.sin(t1)
    y1 = l1 * -np.cos(t1)
    x2 = x1 + l2 * np.sin(t2)
    y2 = y1 + l2 * -np.cos(t2)

    vx1 = l1 * np.cos(t1) * w1
    vy1 = l1 * np.sin(t1) * w1
    vx2 = vx1 + l2 * np.cos(t2) * w2
    vy2 = vy1 + l2 * np.sin(t2) * w2

    e1 = m1 * (vx1*vx1+vy1*vy1) / 2
    e2 = m2 * (vx2*vx2+vy2*vy2) / 2
    p1 = m1 * y1 * 9.8
    p2 = m2 * y2 * 9.8

    e_total = e1 + e2 + p1 + p2

    print(e_total[0])

    plt.subplot(3, 1, 1)
    plt.title('theta1')
    plt.plot(t,t1,linestyle = style, color =color, label = res_tag + '_' + res_index)
    plt.subplot(3, 1, 2)
    plt.title('theta2')
    plt.plot(t,t2, linestyle=style, color=color, label=res_tag + '_' + res_index)
    plt.subplot(3, 1, 3)
    plt.title('Energy')
    plt.plot(t,e_total, linestyle=style, color=color, label= res_tag + '_' + res_index)
 #   plt.subplot(4, 1, 4)
 #   plt.title('phase')
 #   t,x = load(res_index, res_tag, 'phase')
 #   plt.plot(t,x, linestyle=style, color=color, label= res_tag + '_' + res_index)
    plt.legend()
    plt.show()
 #   return
    fig = plt.figure(dpi=144)
    ax = fig.gca()
    ax.set_xlim(-3, 3)
    ax.set_ylim(-3, 3)
    ax.set_aspect("equal")
    ax.grid()

    pendumlum_mark, = ax.plot([], [], "-o", lw=2)
    time_mark = ax.text(0.05, 0.9, '', transform=ax.transAxes)

    def init():
        x = [0.0, x1[0], x2[0]]
        y = [0.0, y1[0], y2[0]]
        pendumlum_mark.set_data(x, y)
        time_mark.set_text('')
        return pendumlum_mark, time_mark

    def update(num):
        x = [0.0, x1[num], x2[num]]
        y = [0.0, y1[num], y2[num]]

        pendumlum_mark.set_data(x, y)
        time_mark.set_text('time = %.1fs' % (num * h))
        return pendumlum_mark, time_mark

    ani = FuncAnimation(fig, update, range(1, len(y1)),interval=h*100, blit=True, init_func=init)
    # ani.save('dp.mp4', writer='ffmpeg')
    ani.save("dp.gif", writer="pillow")
    plt.show()

def generate_batched(res_tag, style):
    index = 0
    while os.path.exists(root_path + str(index)):
        generate(str(index),res_tag,style,plt.get_cmap('Set1')(index))
        index = index + 1

#generate('euler','solid','orange')
#generate('rk2','dotted','b')
#generate('','rk4','solid','r')
#generate_batched('rk4', 'solid')

generate_DPP('','rk4','solid','r')
#generate_DPP('1','rk4','solid','b')
#plt.legend()
#plt.show()